#include "kernel_operator.h"
using namespace AscendC;
// constexpr int32_t BUFFER_NUM = 2; // tensor num for each queue
constexpr int32_t BUFFER_NUM = 1;
class KernelBallQuery {
public:
    __aicore__ inline KernelBallQuery() {}
    __aicore__ inline void Init(GM_ADDR xyz, GM_ADDR center_xyz, GM_ADDR xyz_batch_cnt, GM_ADDR center_xyz_batch_cnt, GM_ADDR idx, float min_radius, float max_radius, int32_t sample_num, uint32_t totalLength, uint32_t tileNum, uint32_t tilingKey, uint32_t B, uint32_t N, uint32_t M)
    {
        // ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        // ……
        // this->blockLength = totalLength / GetBlockNum();
        this->blockLength = totalLength;
        // this->tileNum = tileNum;    // tileNum = 8
        this->tileNum = 1;
        // ASSERT(tileNum != 0 && "tile num can not be zero!");
        this->tileLength = this->blockLength / tileNum / BUFFER_NUM;
        // printf("blockNum: %d\n", GetBlockNum());
        // printf("blockLength: %d, tileNum: %d, tileLength: %d\n", this->blockLength, this->tileNum, this->tileLength);
        // printf("blockIdx: %d\n", GetBlockIdx());
        // printf("B: %d, N: %d, M: %d\n", B, N, M);

        min_radius = min_radius * min_radius;
        max_radius = max_radius * max_radius;

        if (tilingKey == 1) {
            // printf("[INFO] THIS IS TILINGKEY = 1-1-INIT\n");
            // ……
            xyzGm.SetGlobalBuffer((__gm__ DTYPE_XYZ *)xyz, B * N * 3);
            center_xyzGm.SetGlobalBuffer((__gm__ DTYPE_CENTER_XYZ *)center_xyz, B * M * 3);
            idxGm.SetGlobalBuffer((__gm__ DTYPE_IDX *)idx, B * M * sample_num);

            for(int i = 0; i < B; i++) {
                for(int j = 0; j < M; j++) {
                    int cnt = 0;
                    float x = center_xyzGm(i * M * 3 + j * 3);
                    float y = center_xyzGm(i * M * 3 + j * 3 + 1);
                    float z = center_xyzGm(i * M * 3 + j * 3 + 2);
                    for(int k = 0; k < N; k++) {
                        float x1 = xyzGm(i * N * 3 + k * 3);
                        float y1 = xyzGm(i * N * 3 + k * 3 + 1);
                        float z1 = xyzGm(i * N * 3 + k * 3 + 2);
                        float dist = getLen2(x - x1, y - y1, z - z1);
                        if(dist == 0 || (min_radius <= dist && dist < max_radius)) {
                            if(cnt == 0) {
                                for(int _ = 0; _ < sample_num; _++) {
                                    idxGm(i * M * sample_num + j * sample_num + _) = k;
                                }
                            }
                            idxGm(i * M * sample_num + j * sample_num + cnt) = k;
                            cnt++;
                            if(cnt == sample_num) {
                                break;
                            }
                        }
                    }
                }
            }
        }
        if (tilingKey == 2) {
            // printf("[INFO] THIS IS TILINGKEY = 2-2-INIT\n");
            // ……
            xyzGm.SetGlobalBuffer((__gm__ DTYPE_XYZ *)xyz, N * 3);
            center_xyzGm.SetGlobalBuffer((__gm__ DTYPE_CENTER_XYZ *)center_xyz, M * 3);
            idxGm.SetGlobalBuffer((__gm__ DTYPE_IDX *)idx, M * sample_num);
            xyz_batch_cntGm.SetGlobalBuffer((__gm__ DTYPE_XYZ_BATCH_CNT *)xyz_batch_cnt, B);
            center_xyz_batch_cntGm.SetGlobalBuffer((__gm__ DTYPE_CENTER_XYZ_BATCH_CNT *)center_xyz_batch_cnt, B);

            int N = 0;
            int M = 0;
            int alreadyM = 0;
            int alreadyN = 0;
            for(int i = 0; i < B; i++) {
                alreadyM += M;
                alreadyN += N;
                N = xyz_batch_cntGm(i);
                M = center_xyz_batch_cntGm(i);
                // printf("now: %d, N: %d, M: %d\n", i, N, M);
                for(int j = 0; j < M; j++) {
                    int cnt = 0;
                    float x = center_xyzGm(alreadyM * 3 + j * 3);
                    float y = center_xyzGm(alreadyM * 3 + j * 3 + 1);
                    float z = center_xyzGm(alreadyM * 3 + j * 3 + 2);
                    for(int k = 0; k < N; k++) {
                        float x1 = xyzGm(alreadyN * 3 + k * 3);
                        float y1 = xyzGm(alreadyN * 3 + k * 3 + 1);
                        float z1 = xyzGm(alreadyN * 3 + k * 3 + 2);
                        float dist = getLen2(x - x1, y - y1, z - z1);
                        if(dist < max_radius) {
                            if(cnt == 0) {
                                for(int _ = 0; _ < sample_num; _++) {
                                    idxGm(alreadyM * sample_num + j * sample_num + _) = k;
                                }
                            }
                            idxGm(alreadyM * sample_num + j * sample_num + cnt) = k;
                            cnt++;
                            if(cnt == sample_num) {
                                break;
                            }
                        }
                    }
                }
            }
        }
        /*
        输入shape：xyz: (B, N, 3); center_xyz: (B, M, 3)
        输出shape：(B, M, sample_num)
        b = 5, n = 15, m = 20, sample_num = 2
        */
        /*
        xyzGm: 480, center_xyzGm: 480, xyz_batch_cntGm: 480, center_xyz_batch_cntGm: 480, idxGm: 480
        */
        // 输出所有Gm的长度
        // printf("xyzGm: %d, center_xyzGm: %d, xyz_batch_cntGm: %d, center_xyz_batch_cntGm: %d, idxGm: %d\n", xyzGm.GetSize(), center_xyzGm.GetSize(), xyz_batch_cntGm.GetSize(), center_xyz_batch_cntGm.GetSize(), idxGm.GetSize());
        // 输出xyzGm的所有数据
        // printf("--------------xyzGm: --------------\n");
        // for(int i = 0; i < xyzGm.GetSize(); i++) {
        //     float tmp = xyzGm(i); printf("%f, ", tmp); if((i+1)%3==0) printf("\n");
        // }
        // printf("--------------center_xyzGm: --------------\n");
        // for(int i = 0; i < center_xyzGm.GetSize(); i++) {
        //     float tmp = center_xyzGm(i); printf("%f, ", tmp); if((i+1)%3==0) printf("\n");
        // }
        // printf("--------------xyz_batch_cntGm: --------------\n");
        // for(int i = 0; i < xyz_batch_cntGm.GetSize(); i++) {
        //     int tmp = xyz_batch_cntGm(i); printf("%d, ", tmp); if((i+1)%3==0) printf("\n");
        // }
        // printf("--------------center_xyz_batch_cntGm: --------------\n");
        // for(int i = 0; i < center_xyz_batch_cntGm.GetSize(); i++) {
        //     int tmp = center_xyz_batch_cntGm(i); printf("%d, ", tmp); if((i+1)%3==0) printf("\n");
        // }
        // printf("--------------idxGm: --------------\n");
        // for(int i = 0; i < idxGm.GetSize(); i++) {
        //     int tmp = idxGm(i); printf("%d, ", tmp); if((i+1)%3==0) printf("\n");
        // }
        
    }

private:
    __aicore__ inline float getLen2(float x, float y, float z)
    {
        return x * x + y * y + z * z;
    }

private:
        // ……
    TPipe pipe;
    //create queue for input, in this case depth is equal to buffer num
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueXyz;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueCenterXyz;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueXyzBatchCnt;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueCenterXyzBatchCnt;
    //create queue for output, in this case depth is equal to buffer num
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueIdx;
    GlobalTensor<DTYPE_XYZ> xyzGm;
    GlobalTensor<DTYPE_CENTER_XYZ> center_xyzGm;
    GlobalTensor<int32_t> xyz_batch_cntGm;
    GlobalTensor<int32_t> center_xyz_batch_cntGm;
    GlobalTensor<int32_t> idxGm;

    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
};

extern "C" __global__ __aicore__ void ball_query(GM_ADDR xyz, GM_ADDR center_xyz, GM_ADDR xyz_batch_cnt, GM_ADDR center_xyz_batch_cnt, GM_ADDR idx, GM_ADDR workspace, GM_ADDR tiling) {
    if(GetBlockIdx() != 0) {
        // printf("[INFO] THIS IS NOT BLOCKIDX = 0\n");
        return;
    }
    GET_TILING_DATA(tiling_data, tiling);
    // user kernel impl
    KernelBallQuery op;
    uint32_t tilingKey = 1;
    if (TILING_KEY_IS(1)) {
        // printf("[INFO] THIS IS TILINGKEY = 1\n");
        tilingKey = 1;
    } else if (TILING_KEY_IS(2)) {
        // printf("[INFO] THIS IS TILINGKEY = 2\n");
        tilingKey = 2;
    } else {
        // printf("[INFO] NO TILINGKEY?\n");
        tilingKey = 1;
    }
    op.Init(xyz, center_xyz, xyz_batch_cnt, center_xyz_batch_cnt, idx,
    tiling_data.min_radius, tiling_data.max_radius, tiling_data.sample_num,
    tiling_data.totalLength, tiling_data.tileNum, tilingKey, tiling_data.B, tiling_data.N, tiling_data.M);
}
